#!/usr/bin/env python3

import json
import requests
import unicodedata
import argparse
import os


BASE_URL="https://support.kylinos.cn/protalweb/security/cve/info"

LOG_LEVEL=os.environ.get("CVE_LOG_LEVEL","INFO")


class TablePrinter:
    """终端字符表格打印器，支持中文字符"""
    
    def __init__(self, headers=None, border_style='simple'):
        """
        初始化表格打印器
        
        Args:
            headers: 表头列表
            border_style: 边框样式 ('simple', 'grid', 'fancy')
        """
        self.headers = headers or []
        self.rows = []
        self.border_style = border_style
        
        # 边框字符定义
        self.border_chars = {
            'simple': {
                'horizontal': '-',
                'vertical': '|',
                'corner': '+',
                'cross': '+',
                'top_left': '+',
                'top_right': '+',
                'bottom_left': '+',
                'bottom_right': '+',
                'top_cross': '+',
                'bottom_cross': '+',
                'left_cross': '+',
                'right_cross': '+'
            },
            'grid': {
                'horizontal': '-',
                'vertical': '|',
                'corner': '+',
                'cross': '+',
                'top_left': '+',
                'top_right': '+',
                'bottom_left': '+',
                'bottom_right': '+',
                'top_cross': '+',
                'bottom_cross': '+',
                'left_cross': '+',
                'right_cross': '+'
            },
            'fancy': {
                'horizontal': '─',
                'vertical': '│',
                'corner': '┼',
                'cross': '┼',
                'top_left': '┌',
                'top_right': '┐',
                'bottom_left': '└',
                'bottom_right': '┘',
                'top_cross': '┬',
                'bottom_cross': '┴',
                'left_cross': '├',
                'right_cross': '┤'
            }
        }
    
    def get_display_width(self, text):
        """
        获取字符串的显示宽度（考虑中文字符占用2个字符宽度）
        
        Args:
            text: 要计算宽度的字符串
            
        Returns:
            显示宽度
        """
        width = 0
        for char in str(text):
            # 使用unicodedata判断字符类别
            if unicodedata.east_asian_width(char) in ('F', 'W'):
                # 全角字符（包括中文）占2个字符宽度
                width += 2
            else:
                # 半角字符占1个字符宽度
                width += 1
        return width
    
    def pad_string(self, text, width, align='left'):
        """
        填充字符串到指定宽度
        
        Args:
            text: 要填充的字符串
            width: 目标宽度
            align: 对齐方式 ('left', 'center', 'right')
            
        Returns:
            填充后的字符串
        """
        text = str(text)
        text_width = self.get_display_width(text)
        padding = width - text_width
        
        if padding <= 0:
            return text
        
        if align == 'center':
            left_padding = padding // 2
            right_padding = padding - left_padding
            return ' ' * left_padding + text + ' ' * right_padding
        elif align == 'right':
            return ' ' * padding + text
        else:  # left
            return text + ' ' * padding
    
    def add_row(self, row):
        """
        添加一行数据
        
        Args:
            row: 行数据列表
        """
        self.rows.append(row)
    
    def add_rows(self, rows):
        """
        批量添加行数据
        
        Args:
            rows: 行数据列表的列表
        """
        for row in rows:
            self.add_row(row)
    
    def set_headers(self, headers):
        """
        设置表头
        
        Args:
            headers: 表头列表
        """
        self.headers = headers
    
    def calculate_column_widths(self):
        """
        计算每列的最大宽度
        
        Returns:
            每列宽度的列表
        """
        if not self.headers and not self.rows:
            return []
        
        # 确定列数
        max_cols = 0
        if self.headers:
            max_cols = max(max_cols, len(self.headers))
        for row in self.rows:
            max_cols = max(max_cols, len(row))
        
        # 计算每列的最大宽度
        col_widths = [0] * max_cols
        
        # 检查表头宽度
        if self.headers:
            for i, header in enumerate(self.headers):
                if i < len(col_widths):
                    col_widths[i] = max(col_widths[i], self.get_display_width(header))
        
        # 检查每行数据宽度
        for row in self.rows:
            for i, cell in enumerate(row):
                if i < len(col_widths):
                    col_widths[i] = max(col_widths[i], self.get_display_width(cell))
        
        return col_widths
    
    def create_separator_line(self, col_widths, line_type='middle'):
        """
        创建分隔线
        
        Args:
            col_widths: 每列宽度列表
            line_type: 线条类型 ('top', 'middle', 'bottom')
            
        Returns:
            分隔线字符串
        """
        if not col_widths:
            return ""
        
        chars = self.border_chars[self.border_style]
        
        if line_type == 'top':
            left_char = chars['top_left']
            right_char = chars['top_right']
            cross_char = chars['top_cross']
        elif line_type == 'bottom':
            left_char = chars['bottom_left']
            right_char = chars['bottom_right']
            cross_char = chars['bottom_cross']
        else:  # middle
            left_char = chars['left_cross']
            right_char = chars['right_cross']
            cross_char = chars['cross']
        
        # 构建分隔线
        parts = []
        for i, width in enumerate(col_widths):
            if i == 0:
                parts.append(left_char)
            else:
                parts.append(cross_char)
            parts.append(chars['horizontal'] * (width + 2))  # +2 for padding
        
        parts.append(right_char)
        return ''.join(parts)
    
    def create_data_line(self, row_data, col_widths, align='left'):
        """
        创建数据行
        
        Args:
            row_data: 行数据列表
            col_widths: 每列宽度列表
            align: 对齐方式
            
        Returns:
            格式化的数据行字符串
        """
        if not col_widths:
            return ""
        
        chars = self.border_chars[self.border_style]
        parts = []
        
        for i, width in enumerate(col_widths):
            # 添加竖线
            parts.append(chars['vertical'])
            
            # 添加单元格内容
            if i < len(row_data):
                cell_content = self.pad_string(row_data[i], width, align)
            else:
                cell_content = ' ' * width
            
            parts.append(' ' + cell_content + ' ')
        
        # 添加最后的竖线
        parts.append(chars['vertical'])
        
        return ''.join(parts)
    
    def print_table(self, header_align='center', data_align='left'):
        """
        打印表格
        
        Args:
            header_align: 表头对齐方式 ('left', 'center', 'right')
            data_align: 数据对齐方式 ('left', 'center', 'right')
        """
        col_widths = self.calculate_column_widths()
        
        if not col_widths:
            print("表格为空")
            return
        
        # 打印顶部边框
        print(self.create_separator_line(col_widths, 'top'))
        
        # 打印表头
        if self.headers:
            print(self.create_data_line(self.headers, col_widths, header_align))
            print(self.create_separator_line(col_widths, 'middle'))
        
        # 打印数据行
        for i, row in enumerate(self.rows):
            print(self.create_data_line(row, col_widths, data_align))
            
            # 如果不是最后一行且使用grid样式，添加分隔线
            if self.border_style == 'grid' and i < len(self.rows) - 1:
                print(self.create_separator_line(col_widths, 'middle'))
        
        # 打印底部边框
        print(self.create_separator_line(col_widths, 'bottom'))
    
    def get_table_string(self, header_align='center', data_align='left'):
        """
        获取表格字符串（不直接打印）
        
        Args:
            header_align: 表头对齐方式
            data_align: 数据对齐方式
            
        Returns:
            表格字符串
        """
        import io
        import sys
        
        # 重定向stdout来捕获打印内容
        old_stdout = sys.stdout
        sys.stdout = buffer = io.StringIO()
        
        try:
            self.print_table(header_align, data_align)
            result = buffer.getvalue()
        finally:
            sys.stdout = old_stdout
        
        return result


class CveData:
    def __init__(self, cve_id: str, data: dict = {}) -> None:
        self.cve_id = cve_id
        self.data = data
        self._parse_cve_data(data)
    
    def _parse_cve_data(self, data):
        if data is None or data.get("obj") is None:
            raise ValueError(f"None data for {self.cve_id} with {data}")
        product_list = data.get("obj").get("product_component_list")
        # print(product_list)
        if product_list is None:
            raise ValueError(f"None product_component_list for {self.cve_id} with {data}")
        self.product_list = product_list
        self.threat_severity = data.get("obj").get("threat_severity")
        self.mitigation = data.get("obj").get("mitigation")
        self.cve_detail = data.get("obj").get("cve_detail")
        for cvss in data.get("obj").get("cvss_info_list",[]):
            self.cvss_score = cvss.get("score","未知")
            availability_impact = cvss.get("availability_impact","未知")
            if availability_impact == "HIGH":
                self.availability_impact = "高"
            elif availability_impact == "MEDIUM":
                self.availability_impact = "中"
            elif availability_impact == "LOW":
                self.availability_impact = "低"
            else:
                self.availability_impact = availability_impact
            # TODO: 继续添加 attack_complexity user_interaction integrity_impact confidentiality privileges_required 等解析


def fetch_cve_data(cve_id, no_cache=False):
    if not no_cache:
        if LOG_LEVEL == "DEBUG":
            print("fetch from cache")
        # 从用户家目录下的缓存列表读取,即 $HOME/.cache/cve-kylinos/$CVE_ID
        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "cve-kylinos")
        cache_file = os.path.join(cache_dir, cve_id)
        if os.path.exists(cache_file):
            with open(cache_file, "r") as f:
                return json.loads(f.read())
    ua = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"
    headers = {
        "User-Agent": ua,
        "Content-Type": "application/json"
    }
    form={
        "allTitle":cve_id
    }
    response = requests.post(BASE_URL, headers=headers,data=json.dumps(form))
    if LOG_LEVEL == "DEBUG":
        print("fetch from net")
    if response.status_code != 200:
        print(f"Failed to fetch data for {cve_id} with {response.status_code}")
        return None
    if not no_cache:
        # 保存结果到用户家目录下的缓存列表,即 $HOME/.cache/cve-kylinos/$CVE_ID
        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "cve-kylinos")
        os.makedirs(cache_dir, exist_ok=True)
        cache_file = os.path.join(cache_dir, cve_id)
        with open(cache_file, "w") as f:
            f.write(response.text)
    return response.json()


def serialize_product(product_list:list, filter=None, select=None,select_key:list=['os_info','arch','status','component','version','kysa']):
    """
    序列化产品数据
    
    Args:
        product_list: 产品列表
        filter: 过滤函数
        select: 选择函数
        select_key: 选择的键,可选值为 ['os_info','arch','status','component','version','pkg_source_name','kysa','created_at','update_at']
    """
    for product in product_list:
        status_value = product.get("status",0) 
        status = "待确认"
        if status_value == '2':
            status = "影响"
        elif status_value == '4':
            status = "不影响"

        os_info = "未知"
        product_info=product.get("product_info",None)
        if product_info:
            os_info = product_info.get("product_name","未知")

        data = {
            "status": status,
            "pkg_source_name": product.get("packageSourceName","无"),
            "component": product.get("component_name","无"),
            "version": product.get("version","无"),
            "os_info":os_info,
            "arch": product.get("framework","未知"),
            "kysa": product.get("sa_no","无"),
            "created_at": product.get("created_on",""),
            "update_at":product.get("update_on",""),
        }

        if filter and not filter(data):
            continue

        if select:
            data=select(data)
        yield data


def show_product(cve_data:CveData, show_type:str,select_key:list=['os_info','arch','status','component','version','kysa']):
    """
    显示产品信息
    
    Args:
        product: 产品信息
        show_type: 显示类型 ('table', 'text')
        select_key: 选择的键,可选值为 ['os_info','arch','status','component','version','pkg_source_name','kysa','created_at','update_at']
    """
    if show_type.lower() == "table":
        show_table(cve_data,select_key)
    elif show_type.lower() == "json":
        show_json(cve_data)
    elif show_type.lower() == "csv":
        show_csv(cve_data,output="cve-kylinos.csv",select_key=select_key)
    else:
        raise ValueError(f"不支持的显示类型: {show_type}")


def show_json(cve_data:CveData):
    data = serialize_product(cve_data.product_list)
    result = {
        "cve_id":cve_data.cve_id,
        "threat_severity":cve_data.threat_severity,
        "mitigation":cve_data.mitigation,
        "cve_detail":cve_data.cve_detail,
        "cvss_score":cve_data.cvss_score,
        "product_info": list(data)
    }
    print(json.dumps(result,indent=4,ensure_ascii=False))


def show_csv(cve_data:CveData, output:str,select_key:list=['os_info','arch','status','component','version','kysa']):
    import csv

    header = []
    for key in select_key:
        header.append(HEADER_KEY_MAP.get(key,key))
    def select(product):
        data = []
        for key in select_key:
            if key:
                data.append(product.get(key))
        return data
    data = serialize_product(cve_data.product_list,select=select,select_key=select_key)
    csv_data = list(data)
    print(csv_data)
    with open(output, "w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        # TODO: 在表头的前一列添加 CVE 的信息,并且只占一个单元格
        writer.writerow(header)
        writer.writerows(csv_data)


HEADER_KEY_MAP = {
    "os_info":"操作系统",
    "arch":"架构",
    "status":"状态",
    "component":"组件",
    "version":"版本",
    "pkg_source_name":"源码包",
    "kysa":"安全公告",
    "created_at":"创建时间",
    "update_at":"更新时间",
}

def show_table(cve_data:CveData,select_key:list=['os_info','arch','status','component','version','kysa']):    
    """
    打印表格
    
    Args:
        cve_data: Cve数据信息
        select_key: 选择的键,可选值为 ['os_info','arch','status','component','version','pkg_source_name','kysa','created_at','update_at']
    """
    print(f"CVE: {cve_data.cve_id}\n风险等级: {cve_data.threat_severity}\nCVSS评分: {cve_data.cvss_score}\nCVE详情: {cve_data.cve_detail}")
    header = []
    for key in select_key:
        header.append(HEADER_KEY_MAP.get(key,key))
    def select(product):
        data = []
        for key in select_key:
            if key:
                data.append(product.get(key))
        return data
    data = list(serialize_product(cve_data.product_list,select=select))
            # 创建表格实例
    table = TablePrinter(
        headers=header,
        border_style='fancy')
    
    # 设置表头
    table.add_rows(data)
    print(table.get_table_string())


def run():
    args = argparse.ArgumentParser()
    # 修改前：args.add_argument("cve_id", append=True,help="CVE ID")
    # 修改后：
    args.add_argument("cve_id", nargs='+', help="CVE ID（可输入多个）")
    args.add_argument("--no-cache", default=False, action="store_true", help="不使用缓存")
    args.add_argument("--show-type", default="table",choices=['table','json','csv'], help="显示类型（table、json、csv）")
    args.add_argument("--select", nargs='+', help="选择显示的列", 
                    choices=['os_info','arch','status','component','version','pkg_source_name','kysa','created_at','update_at'],
                    default=['os_info','arch','status','kysa'])
    args = args.parse_args()
    
    for cve in args.cve_id:
        data = fetch_cve_data(cve, args.no_cache)
        if data is None:
            print(f"{cve} 不存在")
            continue
        cve_data = CveData(cve_id=cve,data=data)
        show_product(cve_data, show_type=args.show_type,select_key=args.select)


if __name__ == "__main__":
    run()